#!/usr/bin/env python3
"""
基于pdfminer3k的PDF转TXT工具
使用pdfminer3k库将PDF文件转换为纯文本格式

pdfminer3k的优势:
- 更好的中文文本提取支持
- 能处理复杂的PDF布局
- 对扫描版PDF有更好的兼容性
- 支持多种字符编码

作者: AutoNovel项目
日期: 2025-10-17
"""

import os
import sys
import argparse
from pathlib import Path
from typing import Optional, List
from io import StringIO

try:
    from pdfminer.pdfparser import PDFParser
    from pdfminer.pdfdocument import PDFDocument
    from pdfminer.pdfpage import PDFPage
    from pdfminer.pdfinterp import PDFResourceManager, PDFPageInterpreter
    from pdfminer.pdfdevice import PDFDevice
    from pdfminer.layout import LAParams, LTTextBox, LTTextLine, LTChar
    from pdfminer.converter import TextConverter, PDFPageAggregator
except ImportError:
    print("错误: 请先安装pdfminer3k库")
    print("安装命令: uv add pdfminer3k 或 pip install pdfminer3k")
    sys.exit(1)


class PDFMinerConverter:
    """基于pdfminer3k的PDF转换器类"""

    def __init__(self, encoding: str = 'utf-8'):
        """
        初始化转换器

        Args:
            encoding: 输出文本编码格式，默认为utf-8
        """
        self.encoding = encoding

    def extract_text_from_pdf(self, pdf_path: str, clean: bool = True,
                             layout_analysis: bool = True) -> str:
        """
        使用pdfminer3k从PDF提取文本

        Args:
            pdf_path: PDF文件路径
            clean: 是否清理空白字符
            layout_analysis: 是否进行布局分析

        Returns:
            提取的文本内容
        """
        pdf_path = Path(pdf_path)

        if not pdf_path.exists():
            raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")

        if pdf_path.suffix.lower() != '.pdf':
            raise ValueError("输入文件必须是PDF格式")

        print(f"正在使用pdfminer3k解析: {pdf_path}")

        # 设置输出缓存
        output_string = StringIO()

        try:
            # 创建PDF资源管理器
            rsrcmgr = PDFResourceManager()

            # 设置布局参数
            if layout_analysis:
                laparams = LAParams(
                    line_margin=0.5,
                    char_margin=2.0,
                    line_overlap=0.5,
                    word_margin=0.1,
                    boxes_flow=0.5,
                    detect_vertical=True,
                    all_texts=True
                )
            else:
                laparams = LAParams()

            # 创建设备
            device = TextConverter(rsrcmgr, output_string, laparams=laparams, codec=self.encoding)

            # 创建解释器
            interpreter = PDFPageInterpreter(rsrcmgr, device)

            # 打开PDF文件
            with open(pdf_path, 'rb') as fp:
                # 创建PDF解析器
                parser = PDFParser(fp)

                # 创建PDF文档对象
                document = PDFDocument(parser)

                # 检查文档是否加密
                if document.is_encrypted:
                    try:
                        # 尝试解密
                        document._initialize_password('')
                    except:
                        raise ValueError("PDF文件已加密，无法处理")

                # 处理每一页
                pages = list(PDFPage.create_pages(document))
                total_pages = len(pages)

                print(f"PDF共有 {total_pages} 页")

                for i, page in enumerate(pages, 1):
                    print(f"正在处理第 {i}/{total_pages} 页", end='\r')
                    interpreter.process_page(page)

            # 获取提取的文本
            text = output_string.getvalue()

            # 清理资源
            device.close()
            output_string.close()

            print(f"\n文本提取完成，共提取 {len(text)} 个字符")

            if clean and text:
                text = self.clean_text(text)

            return text

        except Exception as e:
            print(f"pdfminer3k提取过程中发生错误: {str(e)}")
            raise

    def extract_text_with_layout(self, pdf_path: str) -> str:
        """
        使用更精确的布局分析提取文本

        Args:
            pdf_path: PDF文件路径

        Returns:
            提取的文本内容
        """
        pdf_path = Path(pdf_path)

        if not pdf_path.exists():
            raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")

        print(f"正在进行布局分析: {pdf_path}")

        text_parts = []

        try:
            # 创建PDF资源管理器
            rsrcmgr = PDFResourceManager()

            # 设置布局参数 - 更精确的设置
            laparams = LAParams(
                line_margin=0.3,
                char_margin=1.5,
                line_overlap=0.3,
                word_margin=0.05,
                boxes_flow=0.2,
                detect_vertical=True,
                all_texts=True,
                heuristics=True
            )

            # 创建布局分析设备
            device = PDFPageAggregator(rsrcmgr, laparams=laparams)
            interpreter = PDFPageInterpreter(rsrcmgr, device)

            # 打开PDF文件
            with open(pdf_path, 'rb') as fp:
                parser = PDFParser(fp)
                document = PDFDocument(parser)

                pages = list(PDFPage.create_pages(document))
                total_pages = len(pages)

                print(f"布局分析处理 {total_pages} 页")

                for i, page in enumerate(pages, 1):
                    print(f"正在分析第 {i}/{total_pages} 页", end='\r')
                    interpreter.process_page(page)

                    # 获取布局结果
                    layout = device.get_result()

                    # 按Y坐标排序文本框
                    text_elements = []
                    for element in layout:
                        if isinstance(element, LTTextBox):
                            text_elements.append((element.y0, element.get_text()))
                        elif isinstance(element, LTTextLine):
                            text_elements.append((element.y0, element.get_text()))

                    # 按从上到下的顺序排序
                    text_elements.sort(key=lambda x: x[0], reverse=True)

                    # 添加页面分隔符
                    if total_pages > 1:
                        text_parts.append(f"\n{'='*60}")
                        text_parts.append(f"第 {i} 页")
                        text_parts.append('='*60)

                    # 添加文本内容
                    for _, text in text_elements:
                        if text.strip():
                            text_parts.append(text.strip())

            # 合并所有文本
            full_text = '\n'.join(text_parts)
            print(f"\n布局分析完成，共提取 {len(full_text)} 个字符")

            return self.clean_text(full_text)

        except Exception as e:
            print(f"布局分析过程中发生错误: {str(e)}")
            # 如果布局分析失败，回退到基本提取
            return self.extract_text_from_pdf(pdf_path)

    def clean_text(self, text: str) -> str:
        """
        清理提取的文本

        Args:
            text: 原始文本

        Returns:
            清理后的文本
        """
        if not text:
            return text

        lines = text.split('\n')
        cleaned_lines = []

        for line in lines:
            # 清理行首行尾空白
            stripped_line = line.strip()

            # 只保留非空行
            if stripped_line:
                cleaned_lines.append(stripped_line)

        return '\n'.join(cleaned_lines)

    def convert_pdf_to_txt(self, pdf_path: str, output_path: Optional[str] = None,
                          method: str = 'auto', clean: bool = True) -> str:
        """
        将PDF文件转换为TXT文件

        Args:
            pdf_path: PDF文件路径
            output_path: 输出TXT文件路径，如果为None则自动生成
            method: 提取方法 ('auto', 'basic', 'layout')
            clean: 是否清理空白字符

        Returns:
            输出文件路径
        """
        pdf_path = Path(pdf_path)

        # 生成输出文件路径
        if output_path is None:
            output_path = pdf_path.with_suffix('.txt')
        else:
            output_path = Path(output_path)

        print(f"转换方法: {method}")
        print(f"输出文件: {output_path}")

        try:
            # 根据方法选择提取方式
            if method == 'auto':
                # 自动选择：先尝试基本方法，如果结果太短则使用布局分析
                text = self.extract_text_from_pdf(pdf_path, clean=False)

                if len(text.strip()) < 50:  # 如果提取的文本太少
                    print("基本方法提取文本较少，切换到布局分析...")
                    text = self.extract_text_with_layout(pdf_path)
                elif clean:
                    text = self.clean_text(text)

            elif method == 'layout':
                text = self.extract_text_with_layout(pdf_path)
            else:  # basic
                text = self.extract_text_from_pdf(pdf_path, clean)

            # 保存到TXT文件
            with open(output_path, 'w', encoding=self.encoding) as f:
                f.write(text)

            print(f"转换完成! 文件已保存到: {output_path}")
            print(f"文件大小: {output_path.stat().st_size} 字节")

            return str(output_path)

        except Exception as e:
            print(f"转换过程中发生错误: {str(e)}")
            raise

    def batch_convert(self, input_dir: str, output_dir: Optional[str] = None,
                     method: str = 'auto', clean: bool = True) -> List[str]:
        """
        批量转换PDF文件

        Args:
            input_dir: 输入目录路径
            output_dir: 输出目录路径，如果为None则使用输入目录
            method: 提取方法
            clean: 是否清理空白字符

        Returns:
            成功转换的文件路径列表
        """
        input_path = Path(input_dir)

        if not input_path.exists():
            raise FileNotFoundError(f"输入目录不存在: {input_path}")

        if output_dir is None:
            output_path = input_path
        else:
            output_path = Path(output_dir)
            output_path.mkdir(exist_ok=True)

        # 查找所有PDF文件
        pdf_files = list(input_path.glob("*.pdf"))

        if not pdf_files:
            print(f"在目录 {input_path} 中未找到PDF文件")
            return []

        print(f"找到 {len(pdf_files)} 个PDF文件")

        converted_files = []

        for pdf_file in pdf_files:
            try:
                output_file = output_path / pdf_file.with_suffix('.txt').name
                result_path = self.convert_pdf_to_txt(
                    pdf_file, output_file, method, clean
                )
                converted_files.append(result_path)
                print("-" * 50)
            except Exception as e:
                print(f"转换文件 {pdf_file} 时出错: {str(e)}")
                continue

        print(f"\n批量转换完成! 成功转换 {len(converted_files)}/{len(pdf_files)} 个文件")
        return converted_files


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="PDF转TXT工具 (pdfminer3k版本) - 使用pdfminer3k库将PDF文件转换为纯文本格式",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 转换单个PDF文件（自动选择方法）
  python pdf_converter_pdfminer.py docs/ch01.pdf

  # 指定输出路径
  python pdf_converter_pdfminer.py docs/ch01.txt docs/ch01_pdfminer.txt

  # 使用布局分析方法
  python pdf_converter_pdfminer.py docs/ch01.txt --method layout

  # 批量转换目录中的所有PDF文件
  python pdf_converter_pdfminer.py --batch docs/

  # 批量转换并指定输出目录和方法
  python pdf_converter_pdfminer.py --batch docs/ output_txt/ --method layout
        """
    )

    parser.add_argument('input', nargs='?', help='输入PDF文件路径或目录路径（用于批量转换）')
    parser.add_argument('output', nargs='?', help='输出TXT文件路径（可选）')
    parser.add_argument('--batch', '-b', action='store_true',
                       help='批量转换模式，转换指定目录中的所有PDF文件')
    parser.add_argument('--output-dir', '-o', help='批量转换时的输出目录（可选）')
    parser.add_argument('--method', '-m', choices=['auto', 'basic', 'layout'],
                       default='auto', help='文本提取方法（默认: auto）')
    parser.add_argument('--encoding', '-e', default='utf-8',
                       help='输出文件编码格式（默认: utf-8）')
    parser.add_argument('--clean', '-c', action='store_true', default=True,
                       help='清理空白字符（默认启用）')
    parser.add_argument('--no-clean', action='store_true',
                       help='不清理空白字符')

    args = parser.parse_args()

    # 处理清理选项
    clean = args.clean and not args.no_clean

    # 创建转换器
    converter = PDFMinerConverter(encoding=args.encoding)

    try:
        if args.batch:
            # 批量转换模式
            if not args.input:
                print("错误: 批量转换模式需要指定输入目录")
                sys.exit(1)

            converted_files = converter.batch_convert(
                args.input,
                args.output_dir,
                args.method,
                clean
            )

            if converted_files:
                print("\n转换成功的文件:")
                for file_path in converted_files:
                    print(f"  - {file_path}")
            else:
                print("没有文件被成功转换")

        else:
            # 单文件转换模式
            if not args.input:
                print("错误: 请指定要转换的PDF文件")
                parser.print_help()
                sys.exit(1)

            output_path = converter.convert_pdf_to_txt(
                args.input,
                args.output,
                args.method,
                clean
            )

            print(f"\n转换完成! 输出文件: {output_path}")

    except Exception as e:
        print(f"错误: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()